﻿using UnityEngine;
using System.Collections.Generic;
using System.Linq;

namespace CWM.Skinn.FFD
{
    public static class LatticeUtility
    {
        public static LatticeDeformer CreateLatticeDeformer(Mesh m_mesh, Transform transform, LatticeDeformer ffdGizmo = null)
        {
            int x = 1, y = 1, z = 1;
            if (ffdGizmo) { x = ffdGizmo.m_x; y = ffdGizmo.m_y; z = ffdGizmo.m_z; }

            var modelRoot = new GameObject("FFD_" + transform.name);
            modelRoot.transform.CopyWorldSpace(transform);

            var rootBone = new GameObject("rootBone");
            rootBone.transform.SetParent(modelRoot.transform, false);
            rootBone.transform.position = Vector3.zero;

            var smrGo = new GameObject(transform.name);
            smrGo.transform.SetParent(modelRoot.transform, false);

            var localBounds = new Bounds();
            localBounds = m_mesh.vertices.CalculateBounds().Transformed(transform);

            rootBone.transform.position = localBounds.center;

            if (ffdGizmo && ffdGizmo.m_editBounds) { localBounds = ffdGizmo.m_bounds; }

            var dividedBoundsSize = localBounds.GetLatticeSize(x, y, z);
            var vis = modelRoot.AddComponent<LatticeDeformer>();
           
            var scaleOffset = Mathf.Clamp(0.1f * dividedBoundsSize.magnitude, 0.0001f, 0.002f);
            var markerMesh = SkinnEx.BoxPrimative(Color.gray, Vector3.one, Vector3.one, scaleOffset);
            vis.markerSize = scaleOffset;

            var fbv = localBounds.Lattice3D(x, y, z).ToArray();
            var bones = new Transform[fbv.Length + 1];

            var markerShader = Shader.Find("Hidden/Internal-Colored");
            var markerMat = SkinnInternalAsset.GetUnlitMaterial ? SkinnInternalAsset.GetUnlitMaterial : new Material(markerShader);

            for (int i = 0; i < bones.Length; i++)
            {
                if (i < 1) { bones[i] = rootBone.transform; continue; }

                bones[i] = new GameObject(i.ToString()).transform;
                bones[i].position = fbv[i - 1];
                bones[i].SetParent(rootBone.transform);
                bones[i].hideFlags = HideFlags.HideInHierarchy;
                var markerSubGO = bones[i].gameObject;
                var markerMf = markerSubGO.gameObject.AddComponent<MeshFilter>();
                markerMf.sharedMesh = markerMesh;
                var markerMR = markerSubGO.gameObject.AddComponent<MeshRenderer>();
                markerMR.sharedMaterial = markerMat;
                //markerMR.sharedMaterial.hideFlags = HideFlags.HideInInspector;
                markerMR.receiveShadows = false;
                markerMR.shadowCastingMode = UnityEngine.Rendering.ShadowCastingMode.Off;
                markerMR.lightProbeUsage = UnityEngine.Rendering.LightProbeUsage.Off;
                markerMR.reflectionProbeUsage = UnityEngine.Rendering.ReflectionProbeUsage.Off;
                var ffdMarker = markerSubGO.gameObject.AddComponent<LatticeMarker>();
                ffdMarker.lattice = vis;
            }

            var smr = smrGo.gameObject.AddComponent<SkinnedMeshRenderer>();

            var newMesh = m_mesh.Clone() as Mesh;
            smr.sharedMesh = newMesh;
            smr.rootBone = rootBone.transform;
            smr.bones = bones;
            smr.sharedMesh.RecalculateBounds();
            smr.SetBounds();

            Matrix4x4[] bindPoses = new Matrix4x4[bones.Length];
            for (int i = 0; i < bones.Length; i++) bindPoses[i] = bones[i].worldToLocalMatrix * smr.localToWorldMatrix;

            smr.sharedMesh.bindposes = bindPoses;

            for (int i = 1; i < bones.Length; i++) FillVisualLattice(bones[i].GetComponent<LatticeMarker>(), bones, dividedBoundsSize);

            FFDWeightsLattice(smr, localBounds);

            smr.quality = SkinQuality.Bone4;
            smr.updateWhenOffscreen = true;

            var renderer = transform.GetComponent<Renderer>();
            smr.sharedMaterials = renderer ? renderer.sharedMaterials : new Material[1];

            vis.orignalMesh = m_mesh;
            vis.dividedBoundsSize = dividedBoundsSize;
            vis.skinnedMeshRender = smr;
            vis.m_x = x; vis.m_y = y; vis.m_z = z;

            vis.m_bounds = localBounds;

            FillMarkers(smr);
            //RemoveUnusedMarkers(smr);
            return vis;
        }

        public static void FFDWeightsLattice(SkinnedMeshRenderer smr, Bounds bounds)
        {
            var vertices = smr.sharedMesh.vertices;
            vertices.Transformed(out vertices, smr.transform.localToWorldMatrix, false);
            var positions = smr.bones.ToVector3Array();
            var bones = smr.bones;

            BoneWeight[] boneWeights = new BoneWeight[vertices.Length];
            for (int i = 0; i < vertices.Length; i++)
            {
                var closestIndex = positions.GetClosestIndex(vertices[i], 0);
                var bone = bones[closestIndex].GetComponent<LatticeMarker>();
                var distances = new List<Vector2>(positions.Length);
                if (bone)
                {
                    distances.Add(new Vector2(closestIndex, Vector3.Distance(vertices[i], positions[closestIndex])));
                    for (int ii = 0; ii < bone.adjancentIndices.Length; ii++)
                        distances.Add(new Vector2(closestIndex, Vector3.Distance(vertices[i], 2f * positions[bone.adjancentIndices[ii]])));
                }
                else for (int ii = 1; ii < positions.Length; ii++) distances.Add(new Vector2(ii, Vector3.Distance(vertices[i], positions[ii])));
                distances = distances.OrderBy(x => x.y).ToList();

                var channelCount = Mathf.Min(distances.Count, bones.Length > 9 ? 3: 4);
                var weights = new Vector2[channelCount];
                for (int ii = 0; ii < channelCount; ii++) weights[ii] = distances[ii];
                var maxDistance = 1e-38f;
                for (int ii = 0; ii < channelCount; ii++) maxDistance = Mathf.Max(maxDistance, weights[ii].y * 1.01f);
                for (int ii = 0; ii < channelCount; ii++) weights[ii].y = 1f - weights[ii].y / maxDistance;
                var totalWeight = 0f;
                for (int ii = 0; ii < channelCount; ii++) totalWeight += weights[ii].y;
                for (int ii = 0; ii < channelCount; ii++) weights[ii].y /= totalWeight;
                boneWeights[i] = new SkinnEx.BoneWeightData() { weights = weights };
            }
            smr.sharedMesh.boneWeights = boneWeights;
        }

        private static void FillMarkers(SkinnedMeshRenderer smr, bool addWeights = false)
        {
            var bones = smr.bones;
            var boneWeights = smr.GetSharedMesh().boneWeights;
            var ffdMarkers = new LatticeMarker[bones.Length];
            for (int i = 0; i < ffdMarkers.Length; i++) if (bones[i]) ffdMarkers[i] = bones[i].GetComponentInChildren<LatticeMarker>();
            for (int i = 0; i < boneWeights.Length; i++)
            {
                for (int ii = 0; ii < 4; ii++)
                {
                    var w = boneWeights[i].GetBoneWeight(ii); if (w <= Mathf.Epsilon) continue;
                    var ffdMarker = ffdMarkers[boneWeights[i].GetBoneIndex(ii)]; if (!ffdMarker) continue;
                    ffdMarker.weightCount++;
                    if (addWeights) ffdMarker.weights.Add(new Vector2(i, w));
                }
            }
        }

        private static void FillVisualLattice(LatticeMarker marker, Transform[] bones, Vector3 size)
        {
            var transform = marker.transform;
            int b0 = -1, b1 = -1, b2 = -1, b3 = -1;
            for (int i = 0; i < bones.Length; i++)
            {
                if (i < 1 || bones[i] == transform) continue;
                if (IsOrtho(transform.transform, bones[i], size)) b0 = i;
            }
            for (int i = 0; i < bones.Length; i++)
            {
                if (i < 1 || bones[i] == transform) continue;
                if (b0 == i) continue;
                if (IsOrtho(transform.transform, bones[i], size)) b1 = i;
            }
            for (int i = 0; i < bones.Length; i++)
            {
                if (i < 1 || bones[i] == transform) continue;
                if (b0 == i || b1 == i) continue;
                if (IsOrtho(transform.transform, bones[i], size)) b2 = i;
            }
            for (int i = 0; i < bones.Length; i++)
            {
                if (i < 1 || bones[i] == transform) continue;
                if (b0 == i || b1 == i || b2 == i) continue;
                if (IsOrtho(transform.transform, bones[i], size)) b3 = i;
            }

            var adjacency = new List<int>(4);
            if (b0 > -1) adjacency.Add(b0); if (b1 > -1) adjacency.Add(b1); if (b2 > -1) adjacency.Add(b2);
            if (bones.Length > 8 && b3 > -1) adjacency.Add(b3);

            var adjancentMarkers = new LatticeMarker[adjacency.Count];
            for (int i = 0; i < adjacency.Count; i++) adjancentMarkers[i] = bones[adjacency[i]].GetComponentInChildren<LatticeMarker>();
            marker.adjancentMarkers = adjancentMarkers;
            marker.adjancentIndices = adjacency.ToArray();
        }

        private static bool IsOrtho(Transform m_transform, Transform checkedTransform, Vector3 size)
        {
            var dis = Vector3.Distance(m_transform.position, checkedTransform.position);
            const float tolerance = 0.001f;
            var testXYZ = false;
            if (dis <= size.y + tolerance)
            {
                if (m_transform.TransformPoint(new Vector3(0f, dis, 0f)) == checkedTransform.position) testXYZ = true;
                else if (m_transform.TransformPoint(new Vector3(0f, -dis, 0f)) == checkedTransform.position) testXYZ = true;
            }
            if (dis <= size.x + tolerance)
            {
                if (m_transform.TransformPoint(new Vector3(dis, 0f, 0f)) == checkedTransform.position) testXYZ = true;
                else if (m_transform.TransformPoint(new Vector3(-dis, 0f, 0f)) == checkedTransform.position) testXYZ = true;
            }
            if (dis <= size.z + tolerance)
            {
                if (m_transform.TransformPoint(new Vector3(0f, 0f, dis)) == checkedTransform.position) testXYZ = true;
                else if (m_transform.TransformPoint(new Vector3(0f, 0f, -dis)) == checkedTransform.position) testXYZ = true;
            }
            return testXYZ;
        }

        private static void RemoveUnusedMarkers(SkinnedMeshRenderer smr)
        {
            var bones = smr.bones;
            for (int i = 0; i < bones.Length; i++)
            {
                var bone = bones[i]; if (!bone) continue;
                var marker = bone.GetComponent<LatticeMarker>(); if (!marker) continue;
                if (marker.weightCount < 1) marker.gameObject.SetActive(false);
            }
        }

        public static void SmoothWeights(SkinnedMeshRenderer smr, float maxWieght = 1f)
        {
            var vertices = smr.sharedMesh.vertices;
            vertices.Transformed(out vertices, smr.transform.localToWorldMatrix, false);
            var bones = smr.bones;
            int[,] adjacency;
            int adjacencyCount = 16;
            if (!smr.sharedMesh.GetAdjacentIndices(out adjacency, adjacencyCount)) return;
            var boneWeights = smr.sharedMesh.boneWeights;
            var smoothedWeights = new BoneWeight[vertices.Length];
            var duplicates = vertices.GetUniqueIndices();
            for (int i = 0; i < vertices.Length; i++)
            {
                if (duplicates[i] != i) { smoothedWeights[i] = smoothedWeights[duplicates[i]]; continue; }
                var distances = new List<Vector3>(new Vector3[bones.Length]);
                for (int ii = 0; ii < adjacencyCount; ii++)
                {
                    var index = adjacency[i, ii]; if (index < 0) continue;
                    var channelWeights = ((SkinnEx.BoneWeightData)boneWeights[index]).weights;
                    if (channelWeights[0].y > maxWieght) continue;
                    for (int iii = 0; iii < channelWeights.Length; iii++)
                    {
                        if (channelWeights[iii].y <= Mathf.Epsilon) continue;
                        int channelIndex = (int)channelWeights[iii].x;
                        var w = distances[channelIndex];
                        w.x = channelIndex;
                        w.y += channelWeights[iii].y;
                        w.z ++;
                        distances[channelIndex] = w;
                    }
                }

                for (int ii = 0; ii < distances.Count; ii++)
                {
                    if (distances[ii].z < 1) continue;
                    var w = distances[ii];
                    w.y /= w.z;
                    distances[ii] = w;
                }

                {
                    var channelWeights = ((SkinnEx.BoneWeightData)boneWeights[i]).weights;
                    for (int iii = 0; iii < channelWeights.Length; iii++)
                    {
                        if (channelWeights[iii].y <= Mathf.Epsilon) continue;
                        int channelIndex = (int)channelWeights[iii].x;
                        var w = distances[channelIndex];
                        w.x = channelIndex;
                        w.y += channelWeights[iii].y;
                        w.z ++;
                        distances[channelIndex] = w;
                    }
                }

                distances = distances.OrderBy(x => x.y).ToList();
                distances.Reverse();

                var totalWeight = 0f;
                var weights = new Vector2[4];
                for (int ii = 0; ii < weights.Length; ii++) { weights[ii] = distances[ii]; totalWeight += distances[ii].y; };
                for (int ii = 0; ii < weights.Length; ii++) weights[ii].y /= totalWeight;
                smoothedWeights[i] = new SkinnEx.BoneWeightData() { weights = weights };
            }

            smr.sharedMesh = smr.sharedMesh.Clone() as Mesh;
            smr.sharedMesh.boneWeights = smoothedWeights;
        }
    }
}